# coding=utf-8

from numpy import *
from scipy.ndimage import filters
from pylab import *
from PIL import Image
from scipy import ndimage

class ImageBaseFunc:
    def jkdenoise(self, im, u_init, tolerance=0.1, tau=0.125, tv_weight=100):
        m,n = im.shape
        U = u_init
        Px = im
        Py = im
        error = 1

        while (error < tolerance):
            Uold = U

            GradUx = roll(U, -1, axis=1)
            GradUy = roll(U, -1, axis=0)

            PxNew = Px + (tau/tv_weight)*GradUx
            PyNew = Py + (tau/tv_weight)*GradUy
            NormNew = maximum(1, sqrt(PxNew**2 + PyNew**2))

            Px = PxNew/NormNew
            Py = PyNew/NormNew

            RxPx = roll(Px, 1, axis=1)
            RyPy = roll(Py, 1, axis=0)

            DivP = (Px - RxPx) + (Py - RyPy)
            U = im + tv_weight * DivP

            error = linalg.norm(U - Uold)/sqrt(n * m)

        return U, im - U

    def jk_compute_harris_response(self, im, sigma=3):
        imx = zeros(im.shape)
        filters.gaussian_filter(im, (sigma, sigma), (0,1), imx)
        imy = zeros(im.shape)
        filters.gaussian_filter(im, (sigma, sigma), (1,0), imy)

        Wxx = filters.gaussian_filter(imx * imx, sigma)
        Wxy = filters.gaussian_filter(imx * imy, sigma)
        Wyy = filters.gaussian_filter(imy * imy, sigma)

        Wdet = Wxx * Wyy - Wxy ** 2
        Wtr = Wxx + Wyy

        return Wdet / Wtr

    def jk_get_harris_points(self, harrisim, min_dist = 10, threshold = 0.1):
        corner_threshold = harrisim.max() * threshold
        harrisim_t = (harrisim > corner_threshold) * 1

        coords = array(harrisim_t.nonzero()).T

        candidate_values = [harrisim[c[0], c[1]] for c in coords]

        index = argsort(candidate_values)

        allowed_locations = zeros(harrisim.shape)
        allowed_locations[min_dist:-min_dist, min_dist:-min_dist] = 1

        filtered_coords = []
        for i in index:
            if allowed_locations[coords[i, 0], coords[i, 1]] == 1:
                filtered_coords.append(coords[i])
                allowed_locations[(coords[i,0]-min_dist):(coords[i, 0] + min_dist),
                        (coords[i, 1] - min_dist) : (coords[i, 1] + min_dist)] = 0

        return filtered_coords

    def jk_plot_harris_points(self, image, filtered_coords):

        figure()
        gray()
        imshow(image)
        plot([p[1] for p in filtered_coords], [p[0] for p in filtered_coords], '*')
        axis('off')
        show()

    def normalize(self, points):
        """归一化"""
        for row in points:
            row /= points[-1]
        return points
    
    def make_homog(self, points):
        """转化为齐次坐标"""
        return np.vstack((points, np.ones((1, points.shape[1]))))

    def H_from_points(self, fp, tp):
        """[ -x1  -y1  -1  0   0    0  x1x1'  y1x1'  x1'] [h1] = 0"""
        """[ 0     0    0 -x1  -y1  -1 x1y1'  y1y1'  y1'] [h2] """
        """[ -x2  -y2  -1  0   0    0  x2x2'  y2x2'  x2'] [h3]"""
        """[ 0     0    0  -x2 -y2  -1 x2y2'  y2y2'  y2'] [h4] """
        """[       ...                               ] [h5...h9] """
        """使用 svd 奇异值分解 找到 H 的最小二乘解"""
        """fp 映射到 tp。点自动归一化"""
        if fp.shape != tp.shape:
            raise RuntimeError('number of points not match')
        
        # 映射起始点
        m = np.mean(fp[:2], axis = 1)
        maxstd = max(np.std(fp[:2], axis = 1)) + 1e-9
        C1 = np.diag([1/maxstd, 1/maxstd, 1])
        C1[0][2] = -m[0]/maxstd
        C1[0][2] = -m[1]/maxstd
        fp = np.dot(C1, fp)

        # 映射对应点
        m = np.mean(tp[:2], axis = 1)
        maxstd = max(np.std(tp[:2], axis = 1)) + 1e-9
        C2 = np.diag([1/maxstd, 1/maxstd, 1])
        C2[0][2] = -m[0]/maxstd
        C2[1][2] = -m[0]/maxstd
        tp = np.dot(C2, tp)

        nbr_correspondences = fp.shape[1]
        A = np.zeros((2*nbr_correspondences, 9))
        for i in range(nbr_correspondences):
            A[2 * i] = [-fp[0][i], -fp[1][i], -1, 0, 0, 0, tp[0][i]*fp[0][i], tp[0][i]*fp[1][i], tp[0][i]]
            A[2*i + 1] = [0, 0, 0, -fp[0][1], -fp[1][i], -1, tp[1][i]*fp[0][i], tp[1][i]*fp[1][i], tp[1][i]]

        U, S, V = np.linalg.svd(A)
        H = V[8].reshape((3, 3))

        H = np.dot(np.linalg.inv(C2), np.dot(H, C1))

        return H / H[2,2]

    def Haffine_from_points(self, fp, tp):
        """ 仿射变换 """
        if fp.shape != tp.shape:
            raise RuntimeError('number of points do not match')
        
        m = np.mean(fp[:2], axis = 1)
        maxstd = max(np.std(fp[:2], axis = 1)) + 1e-9
        C1 = np.diag([1/maxstd, 1/maxstd, 1])
        C1[0][2] = -m[0]/maxstd
        C1[1][2] = -m[1]/maxstd
        fp_cond = np.dot(C1, fp)

        m = np.mean(tp[:2], axis = 1)
        C2 = C1.copy()
        C2[0][2] = -m[0]/maxstd
        C2[1][2] = -m[1]/maxstd
        tp_cond = np.dot(C2, tp)

        A = np.concatenate((fp_cond[:2], tp_cond[:2]), axis = 0)
        U, S, V = np.linalg.svd(A.T)

        tmp = V[:2].T
        B = tmp[:2]
        C = tmp[2:4]

        tmp2 = np.concatenate((np.dot(C, np.linalg.pinv(B)), np.zeros((2, 1))), axis = 1)
        H = np.vstack((tmp2, [0, 0, 1]))

        H = np.dot(np.linalg.inv(C2), np.dot(H, C1))

        return H / H[2,2]

    def image_in_image(self, im1, im2, tp):
        """一个图片放到另一个图片上，使用仿射变换"""

        m, n = im1.shape[:2]
        fp = np.array([[0,m,m,0],[0,0,n,n],[1,1,1,1]])

        H = self.Haffine_from_points(tp, fp)
        im1_t = ndimage.affine_transform(im1, H[:2, :2], (H[0,2], H[1,2]), im2.shape[:2])
        alpha = (im1_t > 0)

        return (1 - alpha)*im2 + alpha * im1_t

    def demo_affine_test(self, imagefile):
        im = np.array(Image.open(imagefile).convert('L'))
        # H = np.array([[1.4, 0.05, -100], [0.05, 1.5, -100], [0, 0, 1]])
        H = np.array([[1.2, -0.1, -100], [-0.1, 1.2, -100], [0, 0, 1]])
        im2 = ndimage.affine_transform(im, H[:2, :2], (H[0, 2], H[1, 2]))

        figure()
        gray()
        imshow(im2)
        show()

    def demo_image_to_image(self, imagefile1, imagefile2):
        """demo 一个图片放到另一个图片上"""
        im1 = np.array(Image.open(imagefile1).convert('L'))
        im2 = np.array(Image.open(imagefile2).convert('L'))

        # tp = np.array([[264,538,540,264],[40,36,605,605], [1, 1, 1, 1]])
        tp = np.array([[264,538,540,264],[40,36,605,605], [1, 1, 1, 1]])

        im3 = self.image_in_image(im1, im2, tp)

        figure()
        gray()
        imshow(im3)
        axis('equal')
        axis('off')
        show()